# 切小小图，切很多份

import os, glob, shutil
import cv2
from xml.dom.minidom import Document
import numpy as np
import fire
import sys
from numpy.random import randint
from utils.conve import *



def generate_xml(floder_data, filename_data, img_path, img_width, img_height, data_defect, save_path):
    '''
    写xml
    floder_data:xml对应的图片的相对于根目录的目录
    filename_data:对应的图片名称
    img_path:对应的图片的绝对路径
    img_width:图片宽
    img_hight:图片高
    data_defect:缺陷名称
    save_path:保存路径
    '''
    doc = Document()
    annotation_node = doc.createElement("annotation")
    doc.appendChild(annotation_node)

    #在annotation节点下添加floder文本节点
    floder_node = doc.createElement("floder")
    annotation_node.appendChild(floder_node)
    floder_node_data = doc.createTextNode(floder_data)#文件夹名
    floder_node.appendChild(floder_node_data)

    #在annotation节点下添加filename文本节点
    filename_node = doc.createElement("filename")
    annotation_node.appendChild(filename_node)
    filename_node_data = doc.createTextNode(filename_data)#文件名
    filename_node.appendChild(filename_node_data)

    #在annotation节点下添加path文本节点
    path_node = doc.createElement("path")
    annotation_node.appendChild(path_node)
    path_node_data = doc.createTextNode(img_path)#路径名
    path_node.appendChild(path_node_data)

    #添加source节点
    source_node = doc.createElement("source")
    annotation_node.appendChild(source_node)

    #在source节点下添加datebase文本节点
    database_node = doc.createElement("database")
    source_node.appendChild(database_node)
    database_node_data = doc.createTextNode("Unknown")
    database_node.appendChild(database_node_data)

    #在annotation添加size节点
    size_node = doc.createElement("size")
    annotation_node.appendChild(size_node)

    #在size节点下添加width文本节点
    width_node = doc.createElement("width")
    size_node.appendChild(width_node)
    width_node_data = doc.createTextNode(img_width)#图像宽
    width_node.appendChild(width_node_data)

    #在size节点下添加height文本节点
    height_node = doc.createElement("height")
    size_node.appendChild(height_node)
    height_node_data = doc.createTextNode(img_height)#图像高
    height_node.appendChild(height_node_data)

    #在size节点下添加depth文本节点
    depth_node = doc.createElement("depth")
    size_node.appendChild(depth_node)
    depth_node_data = doc.createTextNode("1")#图像通道数
    depth_node.appendChild(depth_node_data)

    #在annotation节点下添加segmented文本节点
    segmented_node = doc.createElement("segmented")
    annotation_node.appendChild(segmented_node)
    segmented_node_data = doc.createTextNode("0")
    segmented_node.appendChild(segmented_node_data)

    #在annotation添加object节点
    object_node = doc.createElement("object")
    annotation_node.appendChild(object_node)

    #在object节点下添加name文本节点
    name_node = doc.createElement("name")
    object_node.appendChild(name_node)
    name_node_data = doc.createTextNode('wg_' + data_defect)#缺陷名称
    name_node.appendChild(name_node_data)

    #在object节点下添加pose文本节点
    pose_node = doc.createElement("pose")
    object_node.appendChild(pose_node)
    pose_node_data = doc.createTextNode("Unspecified")
    pose_node.appendChild(pose_node_data)

    #在object节点下添加truncated文本节点
    truncated_node = doc.createElement("truncated")
    object_node.appendChild(truncated_node)
    truncated_node_data = doc.createTextNode("0")
    truncated_node.appendChild(truncated_node_data)

    #在object节点下添加pose文本节点
    difficult_node = doc.createElement("difficult")
    object_node.appendChild(difficult_node)
    difficult_node_data = doc.createTextNode("0")
    difficult_node.appendChild(difficult_node_data)

    #在object节点下添加bndbox节点
    bndbox_node = doc.createElement("bndbox")
    object_node.appendChild(bndbox_node)

    #在bndbox节点下添加xmin文本节点
    xmin_node = doc.createElement("xmin")
    bndbox_node.appendChild(xmin_node)
    xmin_node_data = doc.createTextNode("10")
    xmin_node.appendChild(xmin_node_data)

    #在bndbox节点下添加ymin文本节点
    ymin_node = doc.createElement("ymin")
    bndbox_node.appendChild(ymin_node)
    ymin_node_data = doc.createTextNode("10")
    ymin_node.appendChild(ymin_node_data)

    #在bndbox节点下添加xmax文本节点
    xmax_node = doc.createElement("xmax")
    bndbox_node.appendChild(xmax_node)
    xmax_node_data = doc.createTextNode("20")
    xmax_node.appendChild(xmax_node_data)

    #在bndbox节点下添加ymax文本节点
    ymax_node = doc.createElement("ymax")
    bndbox_node.appendChild(ymax_node)
    ymax_node_data = doc.createTextNode("20")
    ymax_node.appendChild(ymax_node_data)
    
    
    xml_path = filename_data.split(".")[0] + ".xml"
    file_save_path = os.path.join(save_path, xml_path)
    #print("path: ", file_save_path)
    f = open(file_save_path, "w", encoding='utf-8')
    f.write(doc.toprettyxml(indent = "    "))
    f.close()



def grid_cut_row_lines(img, img_name, row_lines, save_path, top_value, down_value, left_value, right_value, conve):
    '''
    切小图，切线上下拓展，图片左右紧缩
    img:读取的图片
    img_name:图片名字
    row_lines:切线位置
    save_path:保存路径
    top_value:上方扩展
    down_value:下方扩展
    left_value:左边紧缩
    right_value:右边紧缩
    conve:前处理方法，在conve.py中
    '''
    if left_value == 0 and right_value == 0:
        img_cut = img[row_lines-top_value:row_lines+down_value, :, :]
    else:
        img_cut = img[max(row_lines-top_value, 0):min(row_lines+down_value, img.shape[0]), left_value:-right_value, :]
    # print(img_cut.shape)
    
    img_cut = cv2.resize(img_cut, (224, 224)) # 选择在此调整大小
    
    if conve:
        for one_conve in conve.split(','):
            print(one_conve)
            print('img file {} in {}'.format(img_name, one_conve))
            img_cut = eval(one_conve)(img_cut)
        print('<==========end==========>')
    
    cv2.imwrite(os.path.join(save_path, img_name), img_cut)



# glob.glob(os.path.join(img_path, '*.jpg'))
def get_cut_pic(img_path, img_save_path, xml_save_path, count_num, mode, preb, conve):
    '''
    切小小图，传递上下左右参数，循环n次生成n张
    img_path:图片的目录
    img_save_path:保存小图的目录
    xml_save_path:保存xml的目录
    count_num:循环次数，即生成份数
    mode:传入pian/chuan（片/串），控制上下左右和中间切线和宽度
    preb:传入bingpian/bubingpian（并/不并），控制xml标签名和图片名
    conve:前处理方法，在conve.py中
    '''
    for file in glob.glob(os.path.join(img_path, '*.jpg')):
        # print('<==========start==========>')
        print('img file {} is processing'.format(os.path.basename(file)))
        img = cv2.imdecode(np.fromfile(file, dtype=np.uint8), 1)
        # if mode == 'chuan': 
        #     hh, ww = img.shape[:2]
        #     vertical = np.sum(img[:, :, 0], axis = 1)
        #     # 计算梯度
        #     vg = np.gradient(vertical)
        #     y = np.where(vertical < 200)[0][0]
        #     # y = np.argmin(vg[750:])
        # elif mode == 'pian': 
        #     hh, ww = img.shape[:2]
        #     vertical = np.sum(img[:, :, 0], axis = 1)
        #     # 计算梯度
        #     vg = np.gradient(vertical)
        #     y = np.where(vertical < 200)[0][-1]
        #     # y = np.argmax(vg[:375])
        hh, ww = img.shape[:2]
        for i in range(count_num):
            if mode == 'chuan':
                top_value = int(np.random.randint(133, 134, 1))
                down_value = int(np.random.randint(133, 134, 1))
                # down_value = y - 750
                left_value = int(np.random.randint(13, 15, 1))
                right_value = 28 - left_value
                cut_line = hh//2
                width = ww
            elif mode == 'pian':
                # top_value = int(np.random.randint(70, 80, 1))
                down_value = int(np.random.randint(133, 134, 1))
                top_value = 270 - down_value
                # down_value = int(np.random.randint(240, 290, 1))
                # left_value = int(np.random.randint(10, 30, 1))
                # right_value = int(np.random.randint(10 ,30, 1))
                left_value = int(np.random.randint(6, 8, 1))
                right_value = 16 - left_value
                cut_line = hh//2
                width = ww
            file_name = preb + '_' + os.path.basename(file).split('.')[0] + '_' + str(i+1) + '.jpg'

            # try:
            grid_cut_row_lines(img, file_name, cut_line, img_save_path, top_value, down_value, left_value, right_value, conve)           
            # except:
            #     print("img cut error happened!")
            
            try:
                generate_xml(img_path, file_name, os.path.join(img_path, file_name), str(width - left_value - right_value), str(top_value + down_value), preb, xml_save_path)
            except:    
                print("xml make error happened!")


def check_dir(dir_path):
    '''
    检测路径，如果有路径提示是否删除，无路径新建路径
    '''
    if os.path.exists(dir_path):
        print('!!!dir {} exists'.format(dir_path))
        comfirm = input("!!!remove dir, yes or mix? [y or m]: ")
        if comfirm == 'y':
            print('!!!{} have been removed!!!'.format(dir_path))
            shutil.rmtree(dir_path)
            os.makedirs(dir_path)
        elif comfirm != 'm':
            print('did not remove!')
            sys.exit(0)
    else:
        os.makedirs(dir_path)



def check_cls(cls):
    '''
    根据读取的文件夹名定义mode和preb
    '''
    if cls == 'pbp':
        mode = 'pian'
        preb = 'bingpian'
    elif cls == 'pbbp':
        mode = 'pian'
        preb = 'bubingpian'
    elif cls == 'cbp':
        mode = 'chuan'
        preb = 'bingpian'
    elif cls == 'cbbp':
        mode = 'chuan'
        preb = 'bubingpian'
    return mode, preb



def main(root_path, save_path, count_nums = [1, 1, 1, 1], conve = False, dir = True):
    '''
    遍历根目录，切小图写xml，保存n份
    root_path:包含几种类型的图片目录的根目录
    save_path:最后保存小小图的目录
    count_nums:对应pbp、pbbp、cbp、cbbp（片并片、片不并片、串并片、串不并片）的生成份数
    conve:选择前处理方式，False为关闭，字符串为调用方法，以','隔开先后调用，具体方法在conve.py中
    dir:是否新建路径，True新建路径
    '''
    img_save_path = os.path.join(save_path, 'img')
    xml_save_path = os.path.join(save_path, 'xml')
    if dir:
        check_dir(img_save_path)
        check_dir(xml_save_path)
    count_count = {}
    # fin_count = {}
    i = 0
    for cls in ['pbp', 'pbbp', 'cbp', 'cbbp']:
        count_count[cls] = count_nums[i]
        i += 1
    for img_fin_path in os.listdir(root_path):
        img_path = os.path.join(root_path, img_fin_path)
        file_count = len(glob.glob(os.path.join(img_path, '*.jpg')))
        # fin_count[img_fin_path] = count_count[img_fin_path] * file_count
        mode, preb = check_cls(img_fin_path)
        get_cut_pic(img_path, img_save_path, xml_save_path, count_count[img_fin_path], mode, preb, conve)
    # print('bp:{}, bbp:{}'.format(fin_count['pbp']+fin_count['cbp'], fin_count['pbbp']+fin_count['cbbp']))
    # for k in fin_count:
    #     print('{}:{}'.format(k, fin_count[k]))


if __name__ == '__main__':
    fire.Fire(main)